from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnableConfig, RunnablePassthrough
from langchain_openai import ChatOpenAI

from env_utils import LOCAL_API_KEY, LOCAL_BASE_URL

prompt = ChatPromptTemplate.from_messages([
    ('system', '{system_message}'),
    MessagesPlaceholder(variable_name="chat_history", optional=True),
    ('human', '{input}')
])

llm = ChatOpenAI(
    model="qwen3-8b",
    temperature=0.8,
    api_key=LOCAL_API_KEY,
    base_url=LOCAL_BASE_URL,
    extra_body={'chat_template_kwargs': {'enable_thinking': False}},
)
chain = prompt | llm

store = {}

#存储历史消息/InMemoryChatMessageHistory存在内存,数据库存储SQLChatMessageHistory,Redis存储RedisChatMessageHistory
def get_session_history(session_id: str):
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]


chain_with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key='input',
    history_messages_key='chat_history'
)

#剪辑内容，生成摘要
def summarize_messages(current_input):
    session_id = current_input['config']['configurable']['session_id']
    if not session_id:
        raise ValueError("必须通过config参数提供session_id")
    chat_history = get_session_history(session_id)
    stored_messages = chat_history.messages
    if len(stored_messages) < 2:
        return {
            "original_messages": stored_messages,
            "summary_message": None
        }
    #剪辑消息
    last_two_messages = stored_messages[-2:]
    messages_to_summarize = stored_messages[:-2]
    summarize_prompt = ChatPromptTemplate.from_messages([
        ('system', '请将一下历史压缩为一条保留关键信息的摘要信息'),
        # ('placeHolder', '{chat_history}'),
        MessagesPlaceholder(variable_name="chat_history", optional=True),
        ('human', '请生成包含上述对话核心内容的摘要，保留重要事实和决策')
    ])
    summarization_chain = summarize_prompt | llm
    summary_message = summarization_chain.invoke({'chat_history': messages_to_summarize})

    return {
        "original_messages": last_two_messages,
        "summary_message": summary_message,
    }

final_chain = RunnablePassthrough.assign(messages_summarized=summarize_messages) | RunnablePassthrough.assign(
    input=lambda x: x['input'],
    chat_history=lambda x: x['messages_summarized']['original_messages'],
    system_message=lambda
        x: f"你是一个动漫专家，请使用二次元的语气尽你所能回答所有问题。摘要：{x['messages_summarized']['summary_message'].content}" if
        x['messages_summarized'].get('summary_message') else '无摘要'
) | chain_with_message_history

result = final_chain.invoke({'input': '你好我是小白', 'config': {'configurable': {'session_id': 'userTest'}}},
                                           config={'configurable': {'session_id': 'userTest'}})
print(result)
result2 = final_chain.invoke({'input': '我的名字叫什么', 'config': {'configurable': {'session_id': 'userTest'}}},
                                           config={'configurable': {'session_id': 'userTest'}})
print(result2)
result3 = final_chain.invoke({'input': '你知道高达SEED吗？', 'config': {'configurable': {'session_id': 'userTest'}}},
                                           config={'configurable': {'session_id': 'userTest'}})
print(result3)